/*
 * PS/2 (Keyboard and Mouse) Driver
 * main.c
 *
 * Copyright (C) 2016 Aleksandar Andrejevic <theflash@sdf.lonestar.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "ps2.h"
#include <irq.h>
#include <device.h>
#include <timer.h>
#include <heap.h>
#include <stdio.h>
#include <semaphore.h>

const char driver_name[] = "ps2";

static dword_t ps2_init(void);
static dword_t ps2_cleanup(void);
static dword_t ps2_read(device_t *device, void *buffer, size_t length, size_t *bytes_read);
static dword_t ps2_write(device_t *device, const void *buffer, size_t length, size_t *bytes_written);
static dword_t ps2_ioctl(device_t *device, dword_t control_code, const void *in_buffer, size_t in_length, void *out_buffer, size_t out_length);

static char_dev_driver_t ps2_driver_block = {
    .init_proc = ps2_init,
    .cleanup_proc = ps2_cleanup,
    .read_proc = ps2_read,
    .write_proc = ps2_write,
    .ioctl_proc = ps2_ioctl
};

static semaphore_t irq_mutex[2];
static ps2_device_t *ps2_devices[2] = { NULL };

static inline bool_t poll_and_read(byte_t *byte)
{
    qword_t end_time = timer_get_milliseconds() + (qword_t)PS2_TIMEOUT;

    while (!(cpu_read_port_byte(PS2_STATUS_PORT) & PS2_OUTPUT_BUFFER_FULL))
    {
        if (timer_get_milliseconds() >= end_time) return FALSE;
    }

    *byte = cpu_read_port_byte(PS2_DATA_PORT);
    return TRUE;
}

static inline bool_t poll_and_write(byte_t byte)
{
    qword_t end_time = timer_get_milliseconds() + (qword_t)PS2_TIMEOUT;

    while (cpu_read_port_byte(PS2_STATUS_PORT) & PS2_INPUT_BUFFER_FULL)
    {
        if (timer_get_milliseconds() >= end_time) return FALSE;
    }

    cpu_write_port_byte(PS2_DATA_PORT, byte);
    return TRUE;
}

static inline byte_t ps2_send_command(byte_t port, byte_t command, int num_params, ...)
{
    int i;
    int retries = PS2_MAX_RETRIES;
    byte_t response;
    va_list args;

    while (retries--)
    {
        cpu_read_port_byte(PS2_DATA_PORT);

        if (port) cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_SEND_TO_SECONDARY);
        if (!poll_and_write(command)) return PS2_DEVICE_RESPONSE_FAILURE;

        va_start(args, num_params);

        for (i = 0; i < num_params; i++)
        {
            if (port) cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_SEND_TO_SECONDARY);
            if (!poll_and_write(va_arg(args, uintptr_t)))
            {
                va_end(args);
                return PS2_DEVICE_RESPONSE_FAILURE;
            }
        }

        va_end(args);

        if (!poll_and_read(&response)) return PS2_DEVICE_RESPONSE_FAILURE;
        if (response != PS2_DEVICE_RESPONSE_RESEND) break;
    }

    return response;
}

static void ps2_irq_handler(registers_t *registers, byte_t int_num)
{
    byte_t port = int_num == PS2_FIRST_IRQ ? 0 : 1;
    release_mutex(&irq_mutex[port]);

    if (!(cpu_read_port_byte(PS2_STATUS_PORT) & PS2_OUTPUT_BUFFER_FULL)) return;
    byte_t buffer = cpu_read_port_byte(PS2_DATA_PORT);

    if (buffer == 0xAA && ps2_devices[port] == NULL)
    {
        ps2_device_type_t type;

        if (ps2_send_command(port, PS2_DEVICE_CMD_DISABLE_SCAN, 0) != PS2_DEVICE_RESPONSE_ACK) return;
        if (ps2_send_command(port, PS2_DEVICE_CMD_IDENTIFY, 0) != PS2_DEVICE_RESPONSE_ACK) return;

        if (poll_and_read(&buffer))
        {
            byte_t second_byte;

            if (poll_and_read(&second_byte))
            {
                if (port == 0) type = PS2_DEVICE_KEYBOARD;
                else type = PS2_DEVICE_KEYBOARD;
            }
            else
            {
                if (port == 0) type = PS2_DEVICE_MOUSE;
                else type = PS2_DEVICE_MOUSE;
            }
        }
        else
        {
            type = PS2_DEVICE_KEYBOARD;
        }

        if (type == PS2_DEVICE_KEYBOARD)
        {
            if (ps2_send_command(port, PS2_DEVICE_CMD_SCANCODE_SET, 1, 1) != PS2_DEVICE_RESPONSE_ACK) return;
        }

        if (ps2_send_command(port, PS2_DEVICE_CMD_ENABLE_SCAN, 0) != PS2_DEVICE_RESPONSE_ACK) return;

        if (type == PS2_DEVICE_KEYBOARD)
        {
            ps2_keyboard_t *keyboard = (ps2_keyboard_t*)malloc(sizeof(ps2_keyboard_t));
            if (keyboard == NULL) return;

            sprintf(keyboard->header.header.name, "PS2Keyboard%d", port);
            keyboard->header.header.driver= &ps2_driver_block;

            keyboard->header.port = port;
            keyboard->header.type = PS2_DEVICE_KEYBOARD;
            lock_init(&keyboard->header.buffer_lock);
            keyboard->header.buffer_start = keyboard->header.buffer_end = 0;

            if (register_char_device(&keyboard->header.header) == ERR_SUCCESS) ps2_devices[port] = &keyboard->header;
            else free(keyboard);
        }
        else if (type == PS2_DEVICE_MOUSE)
        {
            ps2_mouse_t *mouse = (ps2_mouse_t*)malloc(sizeof(ps2_mouse_t));
            if (mouse == NULL) return;

            sprintf(mouse->header.header.name, "PS2Mouse%d", port);
            mouse->header.header.driver= &ps2_driver_block;

            mouse->header.port = port;
            mouse->header.type = PS2_DEVICE_MOUSE;
            lock_init(&mouse->header.buffer_lock);

            if (register_char_device(&mouse->header.header) == ERR_SUCCESS) ps2_devices[port] = &mouse->header;
            else free(mouse);
        }
    }
    else if (ps2_devices[port])
    {
        lock_acquire(&ps2_devices[port]->buffer_lock);

        if (ps2_devices[port]->type == PS2_DEVICE_KEYBOARD)
        {
            ps2_keyboard_t *keyboard = CONTAINER_OF(ps2_devices[port], ps2_keyboard_t, header);

            if (((ps2_devices[port]->buffer_end + 1) % PS2_BUFFER_CAPACITY) != ps2_devices[port]->buffer_start)
            {
                keyboard->buffer[ps2_devices[port]->buffer_end] = buffer;

                ps2_devices[port]->buffer_end++;
                ps2_devices[port]->buffer_end %= PS2_BUFFER_CAPACITY;
            }
        }
        else if (ps2_devices[port]->type == PS2_DEVICE_MOUSE)
        {
            ps2_mouse_t *mouse = CONTAINER_OF(ps2_devices[port], ps2_mouse_t, header);

            mouse->unfinished_packet[mouse->phase++] = buffer;

            if (mouse->phase == 3)
            {
                if (((ps2_devices[port]->buffer_end + 1) % PS2_BUFFER_CAPACITY) != ps2_devices[port]->buffer_start)
                {
                    ps2_mouse_packet_t packet;

                    packet.flags = mouse->unfinished_packet[0] & 0x07;

                    packet.delta_x = (int32_t)((dword_t)mouse->unfinished_packet[1]);
                    if ((mouse->unfinished_packet[0] >> 4) & 1) packet.delta_x = -packet.delta_x;

                    packet.delta_y = (int32_t)((dword_t)mouse->unfinished_packet[2]);
                    if ((mouse->unfinished_packet[0] >> 5) & 1) packet.delta_y = -packet.delta_y;

                    mouse->buffer[ps2_devices[port]->buffer_end] = packet;

                    ps2_devices[port]->buffer_end++;
                    ps2_devices[port]->buffer_end %= PS2_BUFFER_CAPACITY;
                }

                mouse->phase = 0;
            }
        }

        lock_release(&ps2_devices[port]->buffer_lock);
    }
}

static dword_t ps2_init(void)
{
    init_mutex(&irq_mutex[0], 0);
    init_mutex(&irq_mutex[1], 0);

    byte_t working_ports = 3;
    byte_t buffer;

    cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_DISABLE_FIRST_PORT);
    cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_DISABLE_SECOND_PORT);

    while (cpu_read_port_byte(PS2_STATUS_PORT) & PS2_OUTPUT_BUFFER_FULL) cpu_read_port_byte(PS2_DATA_PORT);

    cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_TEST_CONTROLLER);
    if (!poll_and_read(&buffer)) return ERR_HARDWARE;
    if (buffer != 0x55) return ERR_HARDWARE;

    cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_READ_CONFIG);
    byte_t config;
    if (!poll_and_read(&config)) return ERR_HARDWARE;

    cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_WRITE_CONFIG);
    if (!poll_and_write(1 << 2)) return ERR_HARDWARE;

    cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_TEST_FIRST_PORT);
    if (!poll_and_read(&buffer)) working_ports &= ~(1 << 0);
    if (buffer) working_ports &= ~(1 << 0);

    cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_TEST_SECOND_PORT);
    if (!poll_and_read(&buffer)) working_ports &= ~(1 << 1);
    if (buffer) working_ports &= ~(1 << 1);
    if (!working_ports) return ERR_NOTFOUND;

    if (working_ports & (1 << 0))
    {
        if (register_irq_handler(PS2_FIRST_IRQ, &ps2_irq_handler, FALSE) != ERR_SUCCESS)
        {
            working_ports &= ~(1 << 0);
        }
    }

    if (working_ports & (1 << 1))
    {
        if (register_irq_handler(PS2_SECOND_IRQ, &ps2_irq_handler, FALSE) != ERR_SUCCESS)
        {
            working_ports &= ~(1 << 1);
        }
    }

    config |= working_ports;
    cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_WRITE_CONFIG);
    if (!poll_and_write(config)) return ERR_HARDWARE;

    if (working_ports & (1 << 0)) cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_ENABLE_FIRST_PORT);
    if (working_ports & (1 << 1)) cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_ENABLE_SECOND_PORT);

    if (working_ports & (1 << 0)) poll_and_write(PS2_DEVICE_CMD_RESET);

    if (working_ports & (1 << 1))
    {
        cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_SEND_TO_SECONDARY);
        poll_and_write(PS2_DEVICE_CMD_RESET);
    }

    return ERR_SUCCESS;
}

static dword_t ps2_cleanup(void)
{
    int i;

    cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_DISABLE_FIRST_PORT);
    cpu_write_port_byte(PS2_CONTROL_PORT, PS2_CMD_DISABLE_SECOND_PORT);

    unregister_irq_handler(PS2_FIRST_IRQ, &ps2_irq_handler);
    unregister_irq_handler(PS2_SECOND_IRQ, &ps2_irq_handler);

    for (i = 0; i < 2; i++)
    {
        unregister_char_device(&ps2_devices[i]->header);
        free(ps2_devices[i]);
    }

    return ERR_SUCCESS;
}

static dword_t ps2_read(device_t *device, void *buffer, size_t length, size_t *bytes_read)
{
    ps2_device_t *ps2_device = CONTAINER_OF(device, ps2_device_t, header);

    if (ps2_device->type == PS2_DEVICE_KEYBOARD)
    {
        ps2_keyboard_t *keyboard = CONTAINER_OF(ps2_device, ps2_keyboard_t, header);
        byte_t *byte_buffer = (byte_t*)buffer;

        while (length > 0)
        {
            lock_acquire(&ps2_device->buffer_lock);

            while (ps2_device->buffer_start != ps2_device->buffer_end)
            {
                *byte_buffer++ = keyboard->buffer[ps2_device->buffer_start];
                ps2_device->buffer_start++;
                ps2_device->buffer_start %= PS2_BUFFER_CAPACITY;
                if (--length == 0) break;
            }

            lock_release(&ps2_device->buffer_lock);
            if (length > 0) wait_mutex(&irq_mutex[ps2_device->port], NO_TIMEOUT);
        }

        return ERR_SUCCESS;
    }
    else if (ps2_device->type == PS2_DEVICE_MOUSE)
    {
        ps2_mouse_t *mouse = CONTAINER_OF(ps2_device, ps2_mouse_t, header);
        ps2_mouse_packet_t *packet_buffer = (ps2_mouse_packet_t*)buffer;
        length /= sizeof(ps2_mouse_packet_t);

        while (length > 0)
        {
            lock_acquire(&ps2_device->buffer_lock);

            while (ps2_device->buffer_start != ps2_device->buffer_end)
            {
                *packet_buffer++ = mouse->buffer[ps2_device->buffer_start];
                ps2_device->buffer_start++;
                ps2_device->buffer_start %= PS2_BUFFER_CAPACITY;
                if (--length == 0) break;
            }

            lock_release(&ps2_device->buffer_lock);
            if (length > 0) wait_mutex(&irq_mutex[ps2_device->port], NO_TIMEOUT);
        }

        return ERR_SUCCESS;
    }
    else
    {
        return ERR_INVALID;
    }
}

static dword_t ps2_write(device_t *device, const void *buffer, size_t length, size_t *bytes_written)
{
    return ERR_INVALID;
}

static dword_t ps2_ioctl(device_t *device, dword_t control_code, const void *in_buffer, size_t in_length, void *out_buffer, size_t out_length)
{
    ps2_device_t *ps2_device = CONTAINER_OF(device, ps2_device_t, header);

    switch (control_code)
    {
    case IOCTL_CHAR_DEV_CHECK_INPUT:
        if (out_length >= sizeof(size_t))
        {
            lock_acquire(&ps2_device->buffer_lock);

            if (ps2_device->buffer_start <= ps2_device->buffer_end)
            {
                *((size_t*)out_buffer) = ps2_device->buffer_end - ps2_device->buffer_start;
            }
            else
            {
                *((size_t*)out_buffer) = ps2_device->buffer_end + PS2_BUFFER_CAPACITY - ps2_device->buffer_start;
            }

            lock_release(&ps2_device->buffer_lock);
            return ERR_SUCCESS;
        }
        else
        {
            return ERR_SMALLBUF;
        }

    default:
        return ERR_INVALID;
    }
}

dword_t driver_load(const char *parameters)
{
    return register_char_dev_driver(&ps2_driver_block);
}
